# 回溯算法笔记

# 遍历选择, 在递归前做选择, 在递归后撤销选择

# L78 子集
def subSets(nums):
    res = []
    track = []

    def backtrack(start):
        # 前序位置，收集节点值
        res.append(track[:])

        # 遍历所有的选择
        for i in range(start, len(nums)):  # start
            # 前序：做选择
            track.append(nums[i])
            # 递归
            backtrack(i+1)
            # 后序：撤销选择
            track.pop()

    backtrack(0)
    return res

nums = [1, 2, 3]
print("{}的所有不重复子集为: {}".format(nums, subSets(nums)))


# L46 全排列问题
def permute(nums):
    res = []  # 所有全排列结果
    track = []  # 当前选择的结果
    def backtrack():
        # base case
        if len(track) >= len(nums):
            res.append(track[:])
            return

        # 遍历选择
        for num in nums:
            # 前序
            if num in track: continue
            track.append(num)

            backtrack()

            # 后序
            track.pop()

    backtrack()
    return res



nums = [1, 2, 3]
print("{} 的全排列为 {}".format(nums, permute(nums)))

# L51 N皇后问题
def solveNQueens(n):
    board = ["."*n for _ in range(n)]
    # 检查 board[row][col]是否可以放置Q
    def isValid(board, row, col):
        # 检查3个方向上有没有皇后：左上、上、右上
        # 检查左上连线是否有皇后
        i = row-1
        j = col-1
        while i >= 0 and j >= 0:
            if board[i][j] == "Q": return False
            i -= 1
            j -= 1

        # 检查上面连线是否有皇后
        for i in range(len(board)):
            if board[i][col] == "Q": return False

        # 检查右上方连线是否有皇后冲突
        i = row-1
        j = col+1
        while i >= 0 and j < n:
            if board[i][j] == "Q":
                return False
            i -= 1
            j += 1

        return True

    res = []
    def backtrack(board, row):
        # 触发结束条件
        if row == len(board):
            res.append(board[:])
            return

        for col in range(len(board[row])):
            # 排除非法情况
            if not isValid(board, row, col):
                continue

            # 前序 做选择
            str_list = list(board[row])
            str_list[col] = "Q"
            board[row] = "".join(str_list)

            # 回溯
            backtrack(board, row+1)

            # 后续 撤销选择
            str_list = list(board[row])
            str_list[col] = "."
            board[row] = "".join(str_list)

    backtrack(board, 0)

    return res

N = 5
print("{}皇后问题的棋牌排列: {}".format(N, solveNQueens(N)))


